#!/usr/bin/env python3

import copy
from netCDF4 import Dataset
from pprint import pprint
import os
import sys
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../utils')
from utils import cmd_args, cli, parse_config, has_key, get_value, set_value, run, copy_netcdf_file, wrf_version, Version
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../operators')
import wrf_operators as wrf

args = cmd_args()

version = wrf_version(args.wrf_root)
if version >= Version('4.0'):
	cli.error('WRFPLUS 4.0 does not pass tangient and adjoint tests!')
if not version in (Version('3.6.1'), Version('3.8.1'), Version('3.9.1')):
	cli.error('Only WRF 3.6.1, 3.8.1 and 3.9.1 have been tested for FSO application!')

config = parse_config(args.config_json)

if config['wrfvar3']['ob_format'] == 1:
	if not args.prepbufr_root:
		if os.getenv('PREPBUFR_ROOT'):
			args.prepbufr_root = os.getenv('PREPBUFR_ROOT')
		else:
			cli.error('Option --prepbufr-root or environment variable PREPBUFR_ROOT need to be set!')
	args.prepbufr_root = os.path.abspath(args.prepbufr_root)
	if not os.path.isdir(args.prepbufr_root):
		cli.error(f'Directory {args.prepbufr_root} does not exist!')
elif config['wrfvar3']['ob_format'] == 2:
	if not args.littler_root:
		if os.getenv('LITTLER_ROOT'):
			args.littler_root = os.getenv('LITTLER_ROOT')
		else:
			cli.error('Option --littler-root or environment variable LITTLER_ROOT need to be set!')
	args.littler_root = os.path.abspath(args.littler_root)
	if not os.path.isdir(args.littler_root):
		cli.error(f'Directory {args.littler_root} does not exist!')

start_time = config['custom']['start_time']
end_time = config['custom']['end_time']
datetime_fmt = 'YYYY-MM-DD_HH:mm:ss'
start_time_str = start_time.format(datetime_fmt)
end_time_str = end_time.format(datetime_fmt)
spinup_hours = get_value(config['custom'], 'spinup_hours', 6)

def run_wps(config, spinup_hours):
	wps_config = copy.deepcopy(config)
	wps_config['custom']['start_time'] = config['custom']['start_time'].subtract(hours=spinup_hours)
	wps_config['custom']['forecast_hours'] += spinup_hours
	wrf.config_wps(args.work_root, args.wps_root, args.geog_root, wps_config, args)
	wrf.run_wps_geogrid(args.work_root, args.wps_root, wps_config, args)
	wrf.run_wps_ungrib_metgrid(args.work_root, args.wps_root, args.bkg_root, wps_config, args)

def run_spinup(config, spinup_hours):
	spinup_config = copy.deepcopy(config)
	spinup_config['custom']['start_time'] = config['custom']['start_time'].subtract(hours=spinup_hours)
	spinup_config['custom']['forecast_hours'] += spinup_hours
	wrf.config_wrf(args.work_root, args.wrf_root, spinup_config, args, tag='spinup')
	# Generate wrfbdy to the final end time.
	wrf.run_real(args.work_root, args.wrf_root, spinup_config, args, tag='spinup')
	# Only run spinup to the start time.
	spinup_config['custom']['forecast_hours'] = spinup_hours
	spinup_config['custom']['end_time'] = spinup_config['custom']['start_time'].add(hours=spinup_hours)
	wrf.config_wrf(args.work_root, args.wrf_root, spinup_config, args, tag='spinup')
	wrf.run_wrf(args.work_root, args.wrfplus_root, spinup_config, args, tag='spinup')

def run_fcst(config, exp):
	work_dir = f'{args.work_root}/{exp}'
	wrf.config_wrf(args.work_root, args.wrf_root, config, args, tag=exp)
	run(f'cp --remove-destination {args.work_root}/spinup/wrf/wrfbdy_d01 {work_dir}/wrf/wrfbdy_d01_{start_time_str}')
	run(f'ln -sf {work_dir}/wrf/wrfbdy_d01_{start_time_str} {work_dir}/wrf/wrfbdy_d01')
	wrfinput = f'{args.work_root}/spinup/wrf/wrfout_d01_{start_time_str}'
	if not os.path.isfile(wrfinput):
		cli.error(f'File {wrfinput} does not exist! Check spinup run!')
	run(f'cp --remove-destination {wrfinput} {work_dir}/wrf/{os.path.basename(wrfinput).replace("wrfout", "wrfinput")}')
	run(f'ln -sf {work_dir}/wrf/{os.path.basename(wrfinput).replace("wrfout", "wrfinput")} {work_dir}/wrf/wrfinput_d01')
	if exp == 'fa':
		set_value(config, ['wrfvar6', 'orthonorm_gradient'], True)
		wrf.config_wrfda(work_dir, args.wrfda_root, config, args)
		if config['wrfvar3']['ob_format'] == 2:
			wrf.run_wrfda_obsproc(work_dir, args.wrfda_root, config, args, littler_root=args.littler_root)
		wrf.run_wrfda_3dvar(work_dir, args.wrfda_root, config, args)
		wrf.run_wrfda_update_bc(work_dir, args.wrfda_root, config, args, update_lowbc=False)
	# NOTE: Why use wrf.exe in WRFPLUS?
	wrf.run_wrf(args.work_root, args.wrfplus_root, config, args, tag=exp)

def interp_ref(config):
	work_dir = f'{args.work_root}/ref'
	ref_config = copy.deepcopy(config)
	ref_config['custom']['start_time'] = config['custom']['end_time']
	if not args.ref_root:
		cli.warning('Use background as reference.')
		args.ref_root = args.bkg_root
	else:
		ref_config['custom']['background'] = {
			'type': 'gfs',
			'file_pattern': 'gdas.t{{ bkg_start_time.format("HH") }}z.pgrb2.*.f{{ "%03d" % bkg_forecast_hour }}',
			'dir_pattern': 'gdas.{{ bkg_start_time.format("YYYYMMDD") }}/{{ bkg_start_time.format("HH") }}'
		}
	wrf.config_wps(work_dir, args.wps_root, args.geog_root, ref_config, args)
	run(f'ln -sf {args.work_root}/wps/geo_em.d01.nc {args.work_root}/ref/wps')
	wrf.run_wps_ungrib_metgrid(work_dir, args.wps_root, args.ref_root, ref_config, args)
	wrf.config_wrf(work_dir, args.wrf_root, ref_config, args)
	wrf.run_real(work_dir, args.wrf_root, ref_config, args)

def calc_final_sens(a, b, c):
	for var_name in ('U', 'V', 'T', 'P'):
		xa = a.variables[var_name]
		xb = b.variables[var_name]
		if not f'A_{var_name}' in c.variables: c.createVariable(f'A_{var_name}', xa.datatype, xa.dimensions)
		xc = c.variables[f'A_{var_name}']
		xc.setncatts(xa.__dict__)
		xc[:] = 0.0
		# FIXME: Remove boundary effects.
		xc[:] = xa[:] - xb[:]
		if var_name == 'T':
			xc[:] = xc[:] * (9.8 / 3)**2
		elif var_name == 'P':
			xc[:] = xc[:] * (1.0 / 300.0)**2

def calc_fcst_error():
	xt  = Dataset(f'{args.work_root}/ref/wrf/wrfinput_d01_{end_time_str}', 'r')
	xaf = Dataset(f'{args.work_root}/fa/wrf/wrfout_d01_{end_time_str}',    'r')
	xbf = Dataset(f'{args.work_root}/fb/wrf/wrfout_d01_{end_time_str}',    'r')
	if not os.path.isfile(f'{args.work_root}/fa/wrfplus/final_sens_d01') or args.force:
		cli.notice(f'Calculate final sensitivity {args.work_root}/fa/wrfplus/final_sens_d01.')
		run(f'mkdir -p {args.work_root}/fa/wrfplus')
		run(f'cp {args.work_root}/fa/wrf/wrfout_d01_{end_time_str} {args.work_root}/fa/wrfplus/final_sens_d01')
		saf = Dataset(f'{args.work_root}/fa/wrfplus/final_sens_d01', 'r+')
		calc_final_sens(xaf, xt, saf)
		saf.close()
	else:
		run(f'ls -l {args.work_root}/fa/wrfplus/final_sens_d01')
	if not os.path.isfile(f'{args.work_root}/fb/wrfplus/final_sens_d01'):
		cli.notice(f'Calculate final sensitivity {args.work_root}/fb/wrfplus/final_sens_d01.')
		run(f'mkdir -p {args.work_root}/fb/wrfplus')
		run(f'cp {args.work_root}/fb/wrf/wrfout_d01_{end_time_str} {args.work_root}/fb/wrfplus/final_sens_d01')
		sbf = Dataset(f'{args.work_root}/fb/wrfplus/final_sens_d01', 'r+')
		calc_final_sens(xbf, xt, sbf)
		sbf.close()
	else:
		run(f'ls -l {args.work_root}/fb/wrfplus/final_sens_d01')
	xt.close()
	xaf.close()
	xbf.close()

def calc_fcst_sens(config):
	if not os.path.isdir(f'{args.work_root}/sens'): os.mkdir(args.work_root + '/sens')
	config['wrfvar6']['write_lanczos'] = False
	config['wrfvar6']['use_lanczos'] = True
	config['wrfvar6']['eps'] = 1.0e-5
	config['wrfvar17']['adj_sens'] = True
	config['wrfvar17']['sensitivity_option'] = 0
	config['wrfvar17']['analysis_type'] = 'QC-OBS'
	config['time_control']['io_form_auxinput17'] = 2
	config['time_control']['auxinput17_inname'] = './gr01'
	config['time_control']['iofields_filename'] = f'{args.wrfda_root}/var/run/fso.io_config'
	wrf.config_wrfda(args.work_root + '/sens', args.wrfda_root, config, args, args.work_root + '/fa/wrf')
	cli.notice('Add two init_sens_d01 data.')
	os.chdir(args.work_root + '/sens/wrfda/d01')
	run(f'cp {args.work_root}/fa/wrfplus/init_sens_d01_{start_time_str} ad_d01_{start_time_str}')
	sa = Dataset(f'{args.work_root}/fa/wrfplus/init_sens_d01_{start_time_str}', 'r')
	sb = Dataset(f'{args.work_root}/fb/wrfplus/init_sens_d01_{start_time_str}', 'r')
	ad = Dataset(f'ad_d01_{start_time_str}', 'r+')
	for var_name in ('A_U', 'A_V', 'A_T', 'A_W', 'A_PH', 'A_MU', 'A_QVAPOR'):
		xa = sa.variables[var_name]
		xb = sb.variables[var_name]
		xc = ad.variables[var_name]
		xc[:] = xa[:] + xb[:]
	sa.close()
	sb.close()
	ad.close()
	run(f'ln -sf ad_d01_{start_time_str} gr01')
	if config['wrfvar3']['ob_format'] == 2:
		run(f'ln -sf {args.work_root}/fa/wrfda/obsproc/obs_gts_{start_time_str}.3DVAR ob.ascii')
	wrf.run_wrfda_3dvar(args.work_root + '/sens', args.wrfda_root, config, args, wrf_work_dir=args.work_root + '/fa/wrf')
cli.banner(f'Run WPS')
run_wps(config, spinup_hours)

cli.banner(f'Run {spinup_hours} hours spinup')
run_spinup(config, spinup_hours)

cli.banner(f'Run forecast with xb as initial condition')
run_fcst(config, 'fb')

cli.banner(f'Run forecast with xa as initial condition')
run_fcst(config, 'fa')

# Interpolate reference at valid time.
cli.banner(f'Interpolate reference at valid time')
interp_ref(config)

# Calculate forecast error measures.
cli.banner(f'Calculate forecast error measures')
calc_fcst_error()

# Run adjoint model with forecast error.
cli.banner(f'Run adjoint for forecast from background')
wrf.config_wrfplus(args.work_root + '/fb', args.wrfplus_root, config, args)
wrf.run_wrfplus_ad(args.work_root + '/fb', args.wrfplus_root, config, args)
wrf.config_wrfplus(args.work_root + '/fa', args.wrfplus_root, config, args)
wrf.run_wrfplus_ad(args.work_root + '/fa', args.wrfplus_root, config, args)

# Calculate forecast sensitivity.
cli.banner(f'Calculate forecast sensitivity')
calc_fcst_sens(config)

